from math import exp
import numpy as np
import cv2
import os
import imgproc
import matplotlib.pyplot as plt
import json
import time


def load_json(json_path):
    with open(json_path, mode='r', encoding='utf8') as f:
        data = json.load(f)
    return data

def show_img(imgs: np.ndarray, color=False):
    if (len(imgs.shape) == 3 and color) or (len(imgs.shape) == 2 and not color):
        imgs = np.expand_dims(imgs, axis=0)
    for img in imgs:
        plt.figure()
        plt.imshow(img, cmap=None if color else 'gray')

class GaussianTransformer(object):
    def __init__(self, Gaussian_img_size=512, region_threshold=0.3,distanceRatio = 3.34,create=True,gauss_map_path=None):
        """

        :param Gaussian_img_size: 生成高斯分布图的变长
        :param region_threshold: 看saveGaussianHeat的生成结果，白色的线代表这个阈值，它的作用在于这个白色的框会和target box映射
        :param distanceRatio: 生成高斯分布的一个参数，决定高斯分布的方差
        :param create: 实时创建高斯map还是读取图片生成
        """

        self.distanceRatio=distanceRatio
        self.region_threshold = region_threshold
        self.Gaussian_img_size = Gaussian_img_size
        if create:
            self.standardGaussianHeat = self._gen_gaussian_heatmap()

        else:
            self.standardGaussianHeat= cv2.imread(gauss_map_path,0)
        # 归一化
        self.standardGaussianHeat=self.standardGaussianHeat.astype('float')/self.standardGaussianHeat.max()

        _, binary = cv2.threshold(self.standardGaussianHeat, region_threshold , 1, 0)
        np_contours = np.roll(np.array(np.where(binary != 0)), 1, axis=0).transpose().reshape(-1, 2)
        x, y, w, h = cv2.boundingRect(np_contours)
        self.regionbox = np.array([[x, y], [x + w, y], [x + w, y + h], [x, y + h]], dtype=np.int32)
        # print("regionbox", self.regionbox)
    def _gen_gaussian_heatmap(self):
        # 生成高斯分布图
        scaledGaussian = lambda x: exp(-(1 / 2) * (x ** 2))
        heat = np.zeros((self.Gaussian_img_size, self.Gaussian_img_size), np.uint8)
        for i in range(self.Gaussian_img_size):
            for j in range(self.Gaussian_img_size):
                distanceFromCenter = np.linalg.norm(np.array([i - self.Gaussian_img_size / 2, j - self.Gaussian_img_size / 2]))
                distanceFromCenter = self.distanceRatio * distanceFromCenter / (self.Gaussian_img_size / 2)
                scaledGaussianProb = scaledGaussian(distanceFromCenter)
                heat[i, j] = np.clip(scaledGaussianProb * 255, 0, 255)
        return heat
    def add_region_character(self, image, target_bbox):
        """

        :param image: 输入的图像，二维（全黑色或者部分加了高斯图）
        :param target_bbox: 单个的目标框
        :return: 对目标框内画高斯，加到原图内返回
        """
        # 越界检查
        if  not (np.all(target_bbox >= 0) and np.all(target_bbox[:, 0] <= image.shape[1]) \
        and np.all(target_bbox[:, 1] <=image.shape[0])):

            # print('越界，注意检查')
            return image


        # 变换M  regionbox: 高斯图region_threshold为界限截出来的正方形，target_bbox 要变化成的框
        M = cv2.getPerspectiveTransform(np.float32(self.regionbox), np.float32(target_bbox))
        warped = cv2.warpPerspective(self.standardGaussianHeat.copy(), M, (image.shape[1], image.shape[0]))

        image = np.where(warped > image, warped, image)


        return image
    def generate_region(self, image_size, bboxes):
        """

        :param image_size: w,h
        :param bboxes: 按照x，y顺序给出的二维数组，代表图中的框的位置

        :return: 加上相应框高斯之后的新图
        """
        height, width = image_size[1], image_size[0]
        target = np.zeros([height, width], dtype=np.float32)
        for i in range(len(bboxes)):
            character_bbox = np.array(bboxes[i].copy())

            target = self.add_region_character(target, character_bbox)

        return target
    def saveGaussianHeat(self):
        images_folder = os.path.abspath(os.path.dirname(__file__)) + '/image'
        # 生成高斯模糊的图并存储
        cv2.imwrite(os.path.join(images_folder, 'standard.jpg'), self.standardGaussianHeat)

        warped_color = cv2.applyColorMap(self.standardGaussianHeat, cv2.COLORMAP_JET)
        cv2.polylines(warped_color, [np.reshape(self.regionbox, (-1, 1, 2))], True, (255, 255, 255), thickness=1)
        print(os.path.join(images_folder, 'standard_color.jpg'))
        cv2.imwrite(os.path.join(images_folder, 'standard_color.jpg'), warped_color)
        standardGaussianHeat1 = self.standardGaussianHeat.copy()

        standardGaussianHeat1[standardGaussianHeat1 > 0] = 255
        threshhold_guassian = cv2.applyColorMap(standardGaussianHeat1, cv2.COLORMAP_JET)
        cv2.polylines(threshhold_guassian, [np.reshape(self.regionbox, (-1, 1, 2))], True, (255, 255, 255), thickness=1)
        cv2.imwrite(os.path.join(images_folder, 'threshhold_guassian.jpg'), threshhold_guassian)







def setup_logger(log_file_path: str = None):
    import logging
    from colorlog import ColoredFormatter
    logging.basicConfig(filename=log_file_path,
                        format='%(asctime)s %(levelname)-8s %(filename)s[line:%(lineno)d]: %(message)s',
                        # 定义输出log的格式
                        datefmt='%Y-%m-%d %H:%M:%S', )
    """Return a logger with a default ColoredFormatter."""
    formatter = ColoredFormatter(
        "%(asctime)s %(log_color)s%(levelname)-8s %(reset)s %(filename)s[line:%(lineno)d]: %(message)s",
        datefmt='%Y-%m-%d %H:%M:%S',
        reset=True,
        log_colors={
            'DEBUG': 'blue',
            'INFO': 'green',
            'WARNING': 'yellow',
            'ERROR': 'red',
            'CRITICAL': 'red',
        })

    logger = logging.getLogger('PAN')
    handler = logging.StreamHandler()
    handler.setFormatter(formatter)
    logger.addHandler(handler)
    logger.setLevel(logging.DEBUG)
    logger.info('logger init finished')
    return logger


# --exeTime
def exe_time(func):
    # 暂时无用
    def newFunc(*args, **args2):
        t0 = time.time()
        back = func(*args, **args2)
        print("{} cost {:.3f}s".format(func.__name__, time.time() - t0))
        return back

    return newFunc


def save_json(data, json_path):
    # 暂时无用
    with open(json_path, mode='w', encoding='utf8') as f:
        json.dump(data, f, indent=4)


def draw_bbox(img_path, result, color=(255, 0, 0), thickness=2):
    if isinstance(img_path, str):
        img_path = cv2.imread(img_path)
        # img_path = cv2.cvtColor(img_path, cv2.COLOR_BGR2RGB)
    img_path = img_path.copy()
    for point in result:
        point = point.astype(int)
        cv2.line(img_path, tuple(point[0]), tuple(point[1]), color, thickness)
        cv2.line(img_path, tuple(point[1]), tuple(point[2]), color, thickness)
        cv2.line(img_path, tuple(point[2]), tuple(point[3]), color, thickness)
        cv2.line(img_path, tuple(point[3]), tuple(point[0]), color, thickness)
    return img_path



#







if __name__ == '__main__':

    # ########## 测试GaussianTransformer ################
    print(os.getcwd())
    # gaussian = GaussianTransformer(Gaussian_img_size=512, region_threshold=0.3,
    #                                distanceRatio = 3.34,create=False,gauss_map_path='/home/wudeyang/code/CRAFT-pytorch/utils/image/standard.jpg')
    gaussian = GaussianTransformer(Gaussian_img_size=512, region_threshold=0.3,
                                   distanceRatio = 3.34,create=True,
                                   gauss_map_path="/home/wudeyang/code/CRAFT-pytorch/utils/image/standard.jpg")
    # gaussian.saveGaussianHeat() # 有待改进，gaussianheat保存的时候需要整数，但生成的是小数
    bbox1 = np.array([[100, 0], [400, 0], [200, 100], [100, 100]])


    bbox3 = np.array([[500,500], [540,500], [540, 540], [500, 540]])

    image = gaussian.generate_region((700, 600), [bbox1,bbox3])
    target_gaussian_heatmap_color = imgproc.cvt2HeatmapImg(image.copy() / 255)
    # 测试，将原始的框画在生成的热图上
    cv2.polylines(target_gaussian_heatmap_color, [np.reshape(bbox1, (-1, 1, 2))], True, (255, 255, 255), thickness=1)
    cv2.polylines(target_gaussian_heatmap_color, [np.reshape(bbox3, (-1, 1, 2))], True, (255, 255, 255), thickness=1)
    cv2.imwrite("test.jpg", target_gaussian_heatmap_color)
    #################################################################
